// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/api.h>
#include <nebula/compute/api.h>

#include <cstdlib>
#include <iostream>
#include <memory>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

// Demonstrate registering a user-defined Nebula compute function outside of the Nebula
// source tree

namespace cp = ::nebula::compute;

template<typename TYPE,
        typename = typename std::enable_if<nebula::is_number_type<TYPE>::value |
                                           nebula::is_boolean_type<TYPE>::value |
                                           nebula::is_temporal_type<TYPE>::value>::type>
turbo::Result<std::shared_ptr<nebula::Array>> get_array_data_sample(
        const std::vector<typename TYPE::c_type> &values) {
    using ArrowBuilderType = typename nebula::TypeTraits<TYPE>::BuilderType;
    ArrowBuilderType builder;
    TURBO_RETURN_NOT_OK(builder.Reserve(values.size()));
    TURBO_RETURN_NOT_OK(builder.append_values(values));
    return builder.finish();
}

const cp::FunctionDoc func_doc{
        "User-defined-function usage to demonstrate registering an out-of-tree function",
        "returns x + y + z",
        {"x", "y", "z"},
        "UDFOptions"};

turbo::Status sample_function(cp::KernelContext *ctx, const cp::ExecSpan &batch,
                              cp::ExecResult *out) {
    // return x + y + z
    const int64_t *x = batch[0].array.get_values<int64_t>(1);
    const int64_t *y = batch[1].array.get_values<int64_t>(1);
    const int64_t *z = batch[2].array.get_values<int64_t>(1);
    std::cout<<out->array_span_mutable()->length<<std::endl;
    int64_t *out_values = out->array_span_mutable()->get_values<int64_t>(1);
    for (int64_t i = 0; i < batch.length; ++i) {
        *out_values++ = *x++ + *y++ + *z++;
    }
    return turbo::OkStatus();
}

turbo::Status execute() {
    const std::string name = "add_three";
    auto func = std::make_shared<cp::ScalarFunction>(name, cp::Arity::Ternary(), func_doc);
    cp::ScalarKernel kernel({nebula::int64(), nebula::int64(), nebula::int64()},
                            nebula::int64(), sample_function);

    kernel.mem_allocation = cp::MemAllocation::PREALLOCATE;
    kernel.null_handling = cp::NullHandling::INTERSECTION;

    TURBO_RETURN_NOT_OK(func->add_kernel(std::move(kernel)));

    auto registry = cp::get_function_registry();
    TURBO_RETURN_NOT_OK(registry->add_function(std::move(func)));

    TURBO_MOVE_OR_RAISE(auto x, get_array_data_sample<nebula::Int64Type>({1, 2, 3}));
    TURBO_MOVE_OR_RAISE(auto y, get_array_data_sample<nebula::Int64Type>({4, 5, 6}));
    TURBO_MOVE_OR_RAISE(auto z, get_array_data_sample<nebula::Int64Type>({7, 8, 9}));

    TURBO_MOVE_OR_RAISE(auto res, cp::call_function(name, {x, y, z}));
    auto res_array = res.make_array();
    std::cout << "turbo::Result" << std::endl;
    std::cout << res_array->to_string() << std::endl;
    return turbo::OkStatus();
}

int main(int argc, char **argv) {
    auto status = execute();
    if (!status.ok()) {
        std::cerr << "Error occurred : " << status.message() << std::endl;
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}
